/*
 * To change this license header, choose License Headers in Project Properties.
 * To change this template file, choose Tools | Templates
 * and open the template in the editor.
 */
package pt.mleiria.cipher.stream;

import java.util.logging.Logger;
import pt.mleiria.utils.MathUtils;

/**
 *
 * @author manuel
 */
public class RC4 {

    private final static Logger LOGGER = Logger.getLogger(RC4.class.getName());
    /**
     * Look up table
     */
    private final byte[] S = new byte[256];
    /**
     * Key stretched
     */
    private final byte[] T = new byte[256];
    /**
     * Key
     */
    private final byte[] key;
    /**
     * Key length
     */
    private final int keylen;

    public RC4(final byte[] key) {
        this.key = key;
        this.keylen = key.length;
        initialize();
    }
    /**
     * Initialization of the permutation S
     */
    private void initialize(){
        if (keylen < 1 || keylen > 256) {
            throw new IllegalArgumentException(
                    "key must be between 1 and 256 bytes");
        } else {
            for (int i = 0; i < 256; i++) {
                S[i] = (byte) i;
                T[i] = key[i % keylen];
            }
            int j = 0;
            byte tmp;
            for (int i = 0; i < 256; i++) {
                j = (j + S[i] + T[i]) & 0xFF;
                tmp = S[j];
                S[j] = S[i];
                S[i] = tmp;
            }
        }
    }
    /**
     * 
     * @param plaintext
     * @return 
     */
    public byte[] encrypt(final byte[] plaintext) {
        final byte[] ciphertext = new byte[plaintext.length];
        int i = 0, j = 0, k, t;
        byte tmp;
        for (int counter = 0; counter < plaintext.length; counter++) {
            i = (i + 1) & 0xFF;
            j = (j + S[i]) & 0xFF;
            tmp = S[j];
            S[j] = S[i];
            S[i] = tmp;
            t = (S[i] + S[j]) & 0xFF;
            k = S[t];
            ciphertext[counter] = (byte) (plaintext[counter] ^ k);
        }
        return ciphertext;
    }

    public byte[] decrypt(final byte[] ciphertext) {
        return encrypt(ciphertext);
    }
    
    public static void main(String[] args){
        String key = "Key";
        String plaintext = "Plaintext";
        RC4 rc4 = new RC4(key.getBytes());
        byte[] cipher = rc4.encrypt(plaintext.getBytes());
        LOGGER.info(MathUtils.bytesToHex(cipher));
    }
}
